package mw

import (
	"context"

	"github.com/cloudwego/kitex/pkg/endpoint"
	"github.com/cloudwego/kitex/pkg/streaming"
)

/*
const (

	RequestID = "X_REQUEST_ID"

)
*/
type wrappedStream struct {
	streaming.Stream
}

func (w *wrappedStream) RecvMsg(m interface{}) error {
	/*err := w.Stream.RecvMsg(m)
	requestID, _ := metainfo.GetPersistentValue(w.Stream.Context(), RequestID)
	klog.Infof("[AccessLog] requestID: %v, Recv a message: %T(%v)", requestID, m, m)*/
	return w.Stream.RecvMsg(m)
}

func (w *wrappedStream) SendMsg(m interface{}) error {
	// requestID, _ := metainfo.GetPersistentValue(w.Stream.Context(), RequestID)
	// klog.Infof("[AccessLog] requestID: %v, Send a message: %T(%v)", requestID, m, m)
	return w.Stream.SendMsg(m)
}

func newWrappedStream(s streaming.Stream) streaming.Stream {
	return &wrappedStream{s}
}

func AccessLogMiddleware(next endpoint.Endpoint) endpoint.Endpoint {
	return func(ctx context.Context, req, res interface{}) error {
		var Nil interface{} // go 里不能直接 switch nil
		switch Nil {
		case req: // 当前中间件用于客户端，并且是 streaming 方法
			err := next(ctx, req, res)
			// stream 对象要在最终 endpoint return 后才能获取
			if tmp, ok := res.(*streaming.Result); err == nil && ok {
				tmp.Stream = newWrappedStream(tmp.Stream) // 包装 stream 对象
			}
			return err
		case res: // 当前中间件用于服务端，并且是 streaming 方法
			if tmp, ok := req.(*streaming.Args); ok {
				tmp.Stream = newWrappedStream(tmp.Stream) // 包装 stream 对象
			}
		default: // 纯 unary 方法，或 thrift 方法
			// do something else
		}
		return next(ctx, req, res)
	}
}
